#pragma once

#include <vector>
#include <queue>
#include <unordered_map>
#include <functional>
#include <sstream>
#include "TcpLayer.h"
#include <UdpLayer.h>
#include <IcmpLayer.h>
#include "IPv4Layer.h"
#include "PayloadLayer.h"
#include "PacketUtils.h"
#include "SystemUtils.h"
#include <ctime>
#include <fstream>
#include <curl/curl.h>
#include<string>
#include <clickhouse/client.h>
#include <headers/unifiedPacketAttr.h>
#include <headers/HttpStatsCollector.h>

// 最大队列长度（调整以获得最佳速度）
#define MAX_QUEUE_LENGTH 20000
#define WRITE_STAT_LINE(description, counter, measurement) \
		file<< description \
            << ','\
			<< counter \
            << ','\
			<<measurement<< std::endl;

using namespace clickhouse;

void StatsSummary2File(HttpStatsCollector& collector,std::string filename)
	{
        std::ofstream file(filename);
        if (!file.is_open()) {
            std::cerr << "Failed to open file: " << filename << std::endl;
            exit ;
        }
		WRITE_STAT_LINE("Sample time", collector.getGeneralStats().sampleTime, "Seconds");
		WRITE_STAT_LINE("Number of HTTP packets", collector.getGeneralStats().numOfHttpPackets, "Packets");
		WRITE_STAT_LINE("Rate of HTTP packets", collector.getGeneralStats().httpPacketRate.totalRate, "Packets/sec");
		WRITE_STAT_LINE("Number of HTTP flows", collector.getGeneralStats().numOfHttpFlows, "Flows");
		WRITE_STAT_LINE("Rate of HTTP flows", collector.getGeneralStats().httpFlowRate.totalRate, "Flows/sec");
		WRITE_STAT_LINE("Number of HTTP pipelining flows", collector.getGeneralStats().numOfHttpPipeliningFlows, "Flows");
		WRITE_STAT_LINE("Number of HTTP transactions", collector.getGeneralStats().numOfHttpTransactions, "Transactions");
		WRITE_STAT_LINE("Rate of HTTP transactions", collector.getGeneralStats().httpTransactionsRate.totalRate, "Transactions/sec");
		WRITE_STAT_LINE("Total HTTP data", collector.getGeneralStats().amountOfHttpTraffic, "Bytes");
		WRITE_STAT_LINE("Rate of HTTP data", collector.getGeneralStats().httpTrafficRate.totalRate, "Bytes/sec");
		WRITE_STAT_LINE("Average packets per flow", collector.getGeneralStats().averageNumOfPacketsPerFlow, "Packets");
		WRITE_STAT_LINE("Average transactions per flow", collector.getGeneralStats().averageNumOfHttpTransactionsPerFlow, "Transactions");
		WRITE_STAT_LINE("Average data per flow", collector.getGeneralStats().averageAmountOfDataPerFlow, "Bytes");

		WRITE_STAT_LINE("Number of HTTP requests", collector.getRequestStats().numOfMessages, "Requests");
		WRITE_STAT_LINE("Rate of HTTP requests", collector.getRequestStats().messageRate.totalRate, "Requests/sec");
		WRITE_STAT_LINE("Total data in headers", collector.getRequestStats().totalMessageHeaderSize, "Bytes");
		WRITE_STAT_LINE("Average header size", collector.getRequestStats().averageMessageHeaderSize, "Bytes");

		WRITE_STAT_LINE("Number of HTTP responses", collector.getResponseStats().numOfMessages, "Responses");
		WRITE_STAT_LINE("Rate of HTTP responses", collector.getResponseStats().messageRate.totalRate, "Responses/sec");
		WRITE_STAT_LINE("Total data in headers", collector.getResponseStats().totalMessageHeaderSize, "Bytes");
		WRITE_STAT_LINE("Average header size", collector.getResponseStats().averageMessageHeaderSize, "Bytes");
		WRITE_STAT_LINE("Num of responses with content-length", collector.getResponseStats().numOfMessagesWithContentLength, "Responses");
		WRITE_STAT_LINE("Total body size (may be compressed)", collector.getResponseStats().totalContentLengthSize, "Bytes");
		WRITE_STAT_LINE("Average body size", collector.getResponseStats().averageContentLengthSize, "Bytes");
        file.close();
	}


// 包处理类回调：统计、批量存入
class packet_handler
{
private:
    int queue_len; 
    std::string table_prefix;
    Client client;
    std::shared_ptr<ColumnUInt32> col_hash_val;
    std::shared_ptr<ColumnInt32> col_packet_len;//type to be confirmed
    std::shared_ptr<ColumnUInt32> col_tv_sec;       //type to be confirmed
    std::shared_ptr<ColumnUInt32> col_tv_nsec;
    std::shared_ptr<ColumnString> col_src_ip;
    std::shared_ptr<ColumnString> col_dst_ip;
    std::shared_ptr<ColumnUInt8> col_ttl;
    std::shared_ptr<ColumnUInt8> col_tos;
    std::shared_ptr<ColumnUInt16> col_id;
    std::shared_ptr<ColumnUInt16> col_offset;
    std::shared_ptr<ColumnUInt8> col_protocol_ip;
    std::shared_ptr<ColumnUInt8> col_protocol;
    std::shared_ptr<ColumnUInt16> col_len_load;
    std::shared_ptr<ColumnUInt16> col_src_port;
    std::shared_ptr<ColumnUInt16> col_dst_port;
    std::shared_ptr<ColumnUInt32> col_ack_num;
    std::shared_ptr<ColumnUInt32> col_seq_num;
    std::shared_ptr<ColumnUInt16> col_flag;
    std::shared_ptr<ColumnUInt16> col_window;
    std::shared_ptr<ColumnUInt16> col_len_udp;
    std::shared_ptr<ColumnUInt8> col_icmp_type;
    std::shared_ptr<ColumnUInt8> col_icmp_code;
    HttpStatsCollector collector;

public:
    // long long packet_count; // 已处理的包个数
    packet_handler(std::string tb_prefix = ""):table_prefix(tb_prefix),client((ClientOptions().SetHost("localhost")))//构造函数
    {
        col_hash_val    = std::make_shared<ColumnUInt32>();
        col_packet_len  =std::make_shared<ColumnInt32>();
        col_tv_sec      = std::make_shared<ColumnUInt32>();
        col_tv_nsec     =std::make_shared<ColumnUInt32>();
        col_src_ip      =std::make_shared<ColumnString>();
        col_dst_ip      =std::make_shared<ColumnString>();
        col_ttl         =std::make_shared<ColumnUInt8>();
        col_tos         =std::make_shared<ColumnUInt8>();
        col_id          =std::make_shared<ColumnUInt16>();
        col_offset      =std::make_shared<ColumnUInt16>();
        col_protocol_ip =std::make_shared<ColumnUInt8>();
        col_protocol     =std::make_shared<ColumnUInt8>();
        col_len_load     =std::make_shared<ColumnUInt16>();
        col_src_port     =std::make_shared<ColumnUInt16>();
        col_dst_port     =std::make_shared<ColumnUInt16>();
        col_ack_num     =std::make_shared<ColumnUInt32>();
        col_seq_num     =std::make_shared<ColumnUInt32>();
        col_flag       =std::make_shared<ColumnUInt16>();
        col_window     =std::make_shared<ColumnUInt16>();
        col_len_udp    =std::make_shared<ColumnUInt16>();
        col_icmp_type  =std::make_shared<ColumnUInt8>();
        col_icmp_code  =std::make_shared<ColumnUInt8>();
        HttpStatsCollector collector();
    };
    ~packet_handler() {
        if(queue_len!=0) send2db();  // 线程池销毁时执行
    }
    //处理函数
    void handle_packet(pcpp::Packet *packet)
    {
        // packet_count++;
        queue_len++;

        unifiedPacketAttr packet_attr(packet);
        col_hash_val->Append(packet_attr.hash_val);
        col_packet_len->Append(packet_attr.packet_len);
        col_tv_sec->Append(packet_attr.tv_sec);
        col_tv_nsec->Append(packet_attr.tv_nsec);
        col_src_ip->Append(packet_attr.src_ip);
        col_dst_ip->Append(packet_attr.dst_ip);
        col_ttl->Append(packet_attr.ttl);
        col_tos->Append(packet_attr.tos);
        col_id->Append(packet_attr.id);
        col_offset->Append(packet_attr.offset);
        col_protocol_ip->Append(packet_attr.protocol_ip);
        col_protocol->Append(packet_attr.protocol);
        col_len_load->Append(packet_attr.len_load);
        col_src_port->Append(packet_attr.src_port);
        col_dst_port->Append(packet_attr.dst_port);
        col_ack_num->Append(packet_attr.ack_num);
        col_seq_num->Append(packet_attr.seq_num);
        col_flag->Append(packet_attr.flag);
        col_window->Append(packet_attr.window);
        col_len_udp->Append(packet_attr.len_udp);
        col_icmp_type->Append(packet_attr.icmp_type);
        col_icmp_code->Append(packet_attr.icmp_code);
        collector.collectStats(packet);
        
        if (queue_len >= MAX_QUEUE_LENGTH) // 如果达到一定数量，就批量存入db
        {
            send2db();
            flush();
            queue_len = 0;
            StatsSummary2File(collector,"/home/njust/program/example-app/stats.txt");
            collector.clear();
        }
    }
    void flush()
    {
        col_hash_val   ->Clear();
        col_packet_len ->Clear();
        col_tv_sec     ->Clear();
        col_tv_nsec    ->Clear();
        col_src_ip     ->Clear();
        col_dst_ip     ->Clear();
        col_ttl        ->Clear();
        col_tos        ->Clear();
        col_id         ->Clear();
        col_offset     ->Clear();
        col_protocol_ip->Clear();
        col_protocol   ->Clear();
        col_len_load   ->Clear();
        col_src_port   ->Clear();
        col_dst_port   ->Clear();
        col_ack_num    ->Clear();
        col_seq_num    ->Clear();
        col_flag       ->Clear();
        col_window     ->Clear();
        col_len_udp    ->Clear();
        col_icmp_type  ->Clear();
        col_icmp_code  ->Clear();
    };
    void send2db()
    { // 送到数据库
        Block block;
        block.AppendColumn("hashVal",       col_hash_val   );
        block.AppendColumn("packetLen",     col_packet_len );
        block.AppendColumn("tvSec",         col_tv_sec     );
        block.AppendColumn("tvNsec",        col_tv_nsec    );
        block.AppendColumn("srcIp",         col_src_ip     );
        block.AppendColumn("dstIp",         col_dst_ip     );
        block.AppendColumn("ttl",           col_ttl        );
        block.AppendColumn("tos",           col_tos        );
        block.AppendColumn("id",            col_id         );
        block.AppendColumn("offset",        col_offset     );
        block.AppendColumn("protocolIP",    col_protocol_ip);
        block.AppendColumn("protocol",      col_protocol   );
        block.AppendColumn("lenLoad",       col_len_load   );
        block.AppendColumn("srcPort",       col_src_port   );
        block.AppendColumn("dstPort",       col_dst_port   );
        block.AppendColumn("ackNum",        col_ack_num    );
        block.AppendColumn("seqNum",        col_seq_num    );
        block.AppendColumn("flag",          col_flag       );
        block.AppendColumn("window",        col_window     );
        block.AppendColumn("lenUdp",        col_len_udp    );
        block.AppendColumn("icmpType",      col_icmp_type  );
        block.AppendColumn("icmpCode",      col_icmp_code  );

        //其他表
        Block block_src_ip;
        block_src_ip.AppendColumn("srcIp",col_src_ip);
        // //block_ip.AppendColumn("Ip",col_dst_ip);
        block_src_ip.AppendColumn("hashVal",col_hash_val);

        Block block_protocol;
        block_protocol.AppendColumn("protocol",col_protocol);
        block_protocol.AppendColumn("hashVal",col_hash_val);
        try
        {
            client.Insert("default."+table_prefix+"_packets", block);
        }
        catch(const std::exception& e)
        {
            std::cerr <<"Insert failue.retrying... error msg:" <<e.what() << '\n';
            try
            {
                client.ResetConnection();
                client.Insert("default."+table_prefix+"_packets", block);
            }
            catch(const std::exception& e) {
                std::cerr<<"insert block failed!"<<std::endl;
            }
        }
        try
        {
            client.Insert("_src_ip_hash_table", block_src_ip);
        }
        catch(const std::exception& e)
        {
            std::cerr <<"Insert failue.retrying... error msg:" <<e.what() << '\n';
            try
            {
                client.ResetConnection();
                client.Insert("_src_ip_hash_table", block_src_ip);
            }
            catch(const std::exception& e)
            {
                // std::this_thread::sleep_for(std::chrono::milliseconds(5));
                std::cerr <<"Insert failue and drop .error msg:" <<e.what() << '\n';
            }
        }

        try
        {
            client.Insert("_protocol_hash_table", block_protocol);
        }
        catch(const std::exception& e)
        {
            std::cerr <<"Insert failue.retrying... error msg:" <<e.what() << '\n';
            try
            {
                client.ResetConnection();
                client.Insert("_protocol_hash_table", block_protocol);
            }
            catch(const std::exception& e)
            {
                // std::this_thread::sleep_for(std::chrono::milliseconds(5));
                std::cerr <<"Insert failue and drop .error msg:" <<e.what() << '\n';
            }
        }
    };
};